package com.aikg.kgrag_java.service.impl;

import com.aikg.kgrag_java.pojo.dto.DataStream;
import com.aikg.kgrag_java.pojo.entity.ConversationInfo;
import com.aikg.kgrag_java.service.IStreamService;
import com.aikg.kgrag_java.utils.JsonUtils;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.github.benmanes.caffeine.cache.Cache;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.net.http.HttpResponse.BodyHandlers;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.stream.Stream;

/**
 * @ClassName: StreamServiceImpl
 * @Description: 流式服务接口实现类
 * @Author: 李宗雨
 * @Date: 2025/1/6 12:25
 * @Version: 1.0
 **/
@Service
public class StreamServiceImpl implements IStreamService {

    // 保存对话的缓存
    @Resource(name = "conversationCache")
    Cache<String, ConversationInfo> conversationCache;

    private final HttpClient httpClient = HttpClient.newHttpClient();

    // todo: 这里仅使用了一个临时的线程池，后续如果添加其他线程池的功能，可以单独搞一个配置线程池的文件
    private final ExecutorService executorService = Executors.newSingleThreadExecutor();

    public Stream<String> getStreamDataByPost(String url, String jsonPayload) throws IOException, InterruptedException {
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create(url))
                .header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(jsonPayload))
                .build();


        HttpResponse<Stream<String>> response = httpClient.send(request, BodyHandlers.ofLines());

        return response.body();
    }

    public Stream<String> getStreamDataByGet(String url) throws IOException, InterruptedException {
         HttpRequest request = HttpRequest.newBuilder()
                    .uri(URI.create(url))
                    .GET()
                    .build();
        HttpResponse<Stream<String>> response = httpClient.send(request, BodyHandlers.ofLines());

        return response.body();
    }

    /**
     * 流式调用接口，对数据增加conversationId后并返回
     * @param url: 请求接口url地址
     * @param bodyParam: post请求的body数据
     * @param conversationId: 对话的uuid
     * @return 流式地回答
     */
    public SseEmitter stream(String url, Object bodyParam, String conversationId) {
        SseEmitter emitter = new SseEmitter();
        try {
            // 使用ObjectMapper将对象转换为JSON字符串
            ObjectMapper objectMapper = new ObjectMapper();
            String jsonRequest = objectMapper.writeValueAsString(bodyParam);
            executorService.submit(() -> {
                try (Stream<String> stream = this.getStreamDataByPost(url, jsonRequest)) {
                    stream.forEach(data -> {
                        try {
                            DataStream dataStream = JsonUtils.toObj(data.replace("data:", ""), DataStream.class);
                            // 增加id字段并发送数据
                            String message = "{\"conversation_id\":\"" +
                                    conversationId +
                                    "\"" +
                                    "\"text\":\"" +
                                    dataStream.getText() +
                                    "\", \"status\":\"" +
                                    dataStream.getStatus() +
                                    "\"}";
                            // 追加答案到本地缓存
                            if(dataStream.getText() != null){
                                conversationCache.getIfPresent(conversationId).appendAnswer(dataStream.getText());
                            }
                            emitter.send(SseEmitter.event().data(message));
                        } catch (IOException e) {
                            emitter.completeWithError(e);
                        }
                    });
                    emitter.complete();
                } catch (IOException | InterruptedException e) {
                    emitter.completeWithError(e);
                }
            });
        }catch (JsonProcessingException e){
            e.printStackTrace();
        }
        return emitter;
    }

    /**
     * 提供流式返回文本信息的功能
     * @param content 文本信息
     * @param conversationId 对话id
     * @return 流式对象
     */
    @Override
    public SseEmitter streamByContent(String content, String conversationId) {
        SseEmitter emitter = new SseEmitter();
        executorService.submit(() -> {
            Stream<String> stream = content.lines();
            stream.forEach(data -> {
                try {
                    // 增加id字段并发送数据
                    String message = "{\"conversation_id\":\"" + conversationId + "\",\"text\":\"" + data + "\",\"status\":\"init\"}";
                    emitter.send(SseEmitter.event().data(message));
                } catch (IOException e) {
                    emitter.completeWithError(e);
                }
            });
            try {
                // 发送结束字段
                String message = "{\"conversation_id\":\"" + conversationId + "\",\"text\":\"\",\"status\":\"finish\"}";
                emitter.send(SseEmitter.event().data(message));
            } catch (IOException e) {
                e.printStackTrace();
            }
            emitter.complete();
        });
        return emitter;
    }


}

